import torch

class ReinforceModelConfig:
    def __init__(self,
                 vocab_size_dic,
                 heads,
                 total_behavior_num,
                 group_num,
                 interest_num=4,
                 short_time=20,
                 id_embedding_dim=64,
                 simple_embedding_dim=8,
                 layers1=1,
                 layers2=1,
                 drop_out=0.2,
                 device_id = 0,
                 compressed=True,
                 diff1=False,
                 diff2=False,
                 layer = 2,
                 use_cos=True,
                 random_init_q=False,
                 gradient_gather=False,
                 compress_net = "gru"
                 ):
        self.vocab_size_dic = vocab_size_dic
        self.heads = heads
        self.interest_num = interest_num
        self.group_num = group_num
        self.drop_out = drop_out
        self.total_behavior_num = total_behavior_num
        self.short_time = short_time
        self.id_embedding_dim = id_embedding_dim
        self.simple_embedding_dim = simple_embedding_dim
        self.device = torch.device(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
        self.compressed = compressed
        self.layers1 = layers1
        self.layers2 = layers2
        self.diff1 = diff1
        self.diff2 = diff2
        self.layer = layer
        self.use_cos = use_cos
        self.random_init_q = random_init_q
        self.gradient_gather = gradient_gather
        self.compress_net = compress_net
